package live.sidian.database.synchronizer.infrastructure.utils;

import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author sidian
 * @date 2021/3/25 12:56
 */
public class CollUtil extends cn.hutool.core.collection.CollUtil {
    /**
     * 获取键重复的元素集合
     *
     * @param <E> 元素类型
     * @param <K> 键类型
     * @return 重复元素键值与重复元素集合的映射
     */
    public static <E, K> Map<K, List<E>> getDupElements(Collection<E> col, Function<E, K> keyGetter) {
        Map<K, List<E>> map = col.stream().collect(Collectors.groupingBy(keyGetter));
        Map<K, List<E>> result = new HashMap<>();
        map.forEach((k, es) -> {
            if (es.size() > 1) {
                result.put(k, es);
            }
        });
        return result;
    }

    /**
     * 求交集
     *
     * @see cn.hutool.core.collection.CollUtil#intersection(Collection, Collection)
     */
    public static <T, K> Collection<T> intersection(Collection<T> coll1, Collection<T> coll2, Function<T, K> keyGetter) {
        return collOp(coll1, coll2, keyGetter, CollUtil::intersection);
    }

    /**
     * 求差集
     *
     * @see cn.hutool.core.collection.CollUtil#subtract(Collection, Collection)
     */
    public static <T, K> Collection<T> subtract(Collection<T> coll1, Collection<T> coll2, Function<T, K> keyGetter) {
        return collOp(coll1, coll2, keyGetter, CollUtil::subtract);
    }

    public static <T> Collection<T> subtract(Collection<T> coll1, Collection<T> coll2) {
        final Collection<T> result = new LinkedList<>(coll1);
        result.removeAll(coll2);
        return result;
    }


    /**
     * 集合操作
     */
    private static <T, K> Collection<T> collOp(Collection<T> coll1, Collection<T> coll2, Function<T, K> keyGetter,
                                               BiFunction<Collection<K>, Collection<K>, Collection<K>> op) {
        Map<K, T> map1 = coll1.stream().collect(Collectors.toMap(keyGetter, Function.identity()));
        Map<K, T> map2 = coll2.stream().collect(Collectors.toMap(keyGetter, Function.identity()));

        Collection<K> keys = op.apply(map1.keySet(), map2.keySet());

        List<T> res = new LinkedList<>();
        keys.forEach(k -> {
            T t = map1.get(k);
            if (t == null) {
                t = map2.get(k);
            }
            res.add(t);
        });
        return res;
    }

    public static <T, K> Map<K, T> toMap(Collection<T> col, Function<T, K> keyGetter) {
        return col.stream().collect(Collectors.toMap(keyGetter, Function.identity()));
    }

    /**
     * 批量获取集合中元素
     *
     * @param coll      集合
     * @param consumer  消费一次批量获取的元素
     * @param batchSize 一次批量获取个数
     * @param <T>       元素类型
     */
    public static <T> void batchGet(Collection<T> coll, java.util.function.Consumer<Collection<T>> consumer, long batchSize) {
        batchSize = batchSize == 0 ? 1 : batchSize;
        List<T> batch = new LinkedList<>();
        for (T element : coll) {
            batch.add(element);
            if (batch.size() == batchSize) {
                consumer.accept(batch);
                batch = new LinkedList<>();
            }
        }
        if (!batch.isEmpty()) {
            consumer.accept(batch);
        }
    }

    /**
     * 批量获取集合中元素
     *
     * @param coll     集合
     * @param consumer 消费一次批量获取的元素
     * @param <T>      元素类型
     */
    public static <T> void batchGet(Collection<T> coll, java.util.function.Consumer<Collection<T>> consumer) {
        batchGet(coll, consumer, 100);
    }
}
